import pandas as pd
from sqlalchemy import create_engine, text
import tkinter as tk
from tkinter import ttk, messagebox, simpledialog
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib import rcParams
import datetime
from tkinter import *
import strategies

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 数据库配置
DB_CONFIG = {
    'host': '127.0.0.1',
    'port': 3306,
    'user': 'root',
    'password': '',
    'database': 'stocks',
    'charset': 'utf8mb4'
}


def fetch_stocks_info(engine):
    """获取stocks_info表所有数据，并按symbol列排序"""
    query = text("SELECT * FROM stocks_info ORDER BY symbol")
    df = pd.read_sql(query, engine)
    # 添加序号列，从1开始
    df.insert(0, '序号', range(1, len(df) + 1))
    return df


def fetch_daily_data(engine, ts_code, start_date=None, end_date=None):
    """获取指定股票的日线数据，并按交易日降序排序"""
    if start_date and end_date:
        query = text("SELECT * FROM stocks_daily WHERE ts_code = :ts_code AND trade_date BETWEEN :start_date AND :end_date ORDER BY trade_date DESC")
        return pd.read_sql(query, engine, params={'ts_code': ts_code, 'start_date': start_date, 'end_date': end_date})
    else:
        query = text("SELECT * FROM stocks_daily WHERE ts_code = :ts_code ORDER BY trade_date DESC")
        return pd.read_sql(query, engine, params={'ts_code': ts_code})


def fetch_date_range(engine):
    """获取数据库中股票数据的日期范围"""
    query = text("SELECT MIN(trade_date) as min_date, MAX(trade_date) as max_date FROM stocks_daily")
    result = pd.read_sql(query, engine)
    if not result.empty:
        return result.iloc[0]['min_date'], result.iloc[0]['max_date']
    return None, None


def select_stock_dialog(parent, engine):
    """选择股票的对话框"""
    # 获取所有股票信息
    stocks_info = fetch_stocks_info(engine)
    
    # 创建对话框
    dialog = tk.Toplevel(parent)
    dialog.title("选择股票")
    dialog.geometry("400x300")
    
    # 创建Treeview组件
    tree = ttk.Treeview(dialog, show="headings")
    tree["columns"] = ['序号', '代码', '名称']
    
    # 设置列格式
    tree.column('序号', width=60, anchor='center')
    tree.column('代码', width=100, anchor='center')
    tree.column('名称', width=100, anchor='center')
    
    tree.heading('序号', text='序号', anchor='center')
    tree.heading('代码', text='股票代码', anchor='center')
    tree.heading('名称', text='股票名称', anchor='center')
    
    # 添加数据
    for i, row in stocks_info.iterrows():
        tree.insert("", tk.END, values=(row['序号'], row['symbol'], row['name']), tags=(row['ts_code'],))
    
    # 滚动条
    vsb = ttk.Scrollbar(dialog, orient="vertical", command=tree.yview)
    tree.configure(yscrollcommand=vsb.set)
    
    # 布局
    tree.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=10, pady=10)
    vsb.pack(side=tk.RIGHT, fill=tk.Y, padx=(0,10), pady=10)
    
    # 选中的股票
    selected_stock = None
    
    def on_select():
        nonlocal selected_stock
        item = tree.selection()[0]
        ts_code = tree.item(item, 'tags')[0]
        selected_stock = (ts_code, tree.item.item(item, 'values')[2])  # (ts_code, name)
        dialog.destroy()
    
    # 选择按钮
    select_btn = ttk.Button(dialog, text="选择", command=on_select)
    select_btn.pack(side=tk.BOTTOM, pady=10)
    
    # 等待对话框关闭
    dialog.wait_window()
    
    return selected_stock


def select_date_range_dialog(parent, min_date, max_date, stock_name=None):
    """选择日期范围的对话框"""
    # 创建对话框
    dialog = tk.Toplevel(parent)
    if stock_name:
        dialog.title(f"选择{stock_name}回测日期范围")
    else:
        dialog.title("选择回测日期范围")
    dialog.geometry("300x200")
    
    # 转换日期格式
    min_date = pd.to_datetime(min_date)
    max_date = pd.to_datetime(max_date)
    
    # 默认日期范围（最近一年）
    default_end_date = max_date
    default_start_date = max_date - pd.DateOffset(years=1)
    
    # 确保默认日期在有效范围内
    if default_start_date < min_date:
        default_start_date = min_date
    
    # 日期选择变量
    start_date_var = tk.StringVar(value=default_start_date.strftime("%Y-%m-%d"))
    end_date_var = tk.StringVar(value=default_end_date.strftime("%Y-%m-%d"))
    
    # 开始日期
    ttk.Label(dialog, text="开始日期:").grid(row=0, column=0, padx=10, pady=10, sticky='w')
    start_entry = ttk.Entry(dialog, textvariable=start_date_var)
    start_entry.grid(row=0, column=1, padx=10, pady=10)
    
    # 结束日期
    ttk.Label(dialog, text="结束日期:").grid(row=1, column=0, padx=10, pady=10, sticky='w')
    end_entry = ttk.Entry(dialog, textvariable=end_date_var)
    end_entry.grid(row=1, column=1, padx=10, pady=10)
    
    # 日期范围提示
    date_range_label = ttk.Label(dialog, text=f"有效日期范围: {min_date.strftime('%Y-%m-%d')} 至 {max_date.strftime('%Y-%m-%d')}")
    date_range_label.grid(row=2, column=0, columnspan=2, padx=10, pady=5)
    
    # 选中的日期范围
    selected_dates = None
    
    def on_confirm():
        nonlocal selected_dates
        try:
            start_date = pd.to_datetime(start_date_var.get())
            end_date = pd.to_datetime(end_date_var.get())
            
            # 验证日期范围
            if start_date < min_date or end_date > max_date:
                messagebox.showerror("错误", f"日期范围必须在 {min_date.strftime('%Y-%m-%d')} 至 {max_date.strftime('%Y-%m-%d')} 之间")
                return
            
            if start_date > end_date:
                messagebox.showerror("错误", "开始日期不能晚于结束日期")
                return
            
            selected_dates = (start_date.strftime("%Y-%m-%d"), end_date.strftime("%Y-%m-%d"))
            dialog.destroy()
            
        except ValueError:
            messagebox.showerror("错误", "请输入有效的日期格式 (YYYY-MM-DD)")
    
    # 确认按钮
    confirm_btn = ttk.Button(dialog, text="确认", command=on_confirm)
    confirm_btn.grid(row=3, column=0, columnspan=2, pady=10)
    
    # 等待对话框关闭
    dialog.wait_window()
    
    return selected_dates


def select_strategy_dialog(parent):
    """选择交易策略的对话框"""
    # 创建对话框
    dialog = tk.Toplevel(parent)
    dialog.title("选择交易策略")
    dialog.geometry("300x200")
    
    # 策略选择变量
    strategy_var = tk.StringVar()
    
    # 获取可用策略列表
    available_strategies = strategies.get_available_strategies()
    strategies_list = [(name, code) for name, code, func in available_strategies]
    
    # 添加策略选项
    for text, value in strategies_list:
        ttk.Radiobutton(dialog, text=text, variable=strategy_var, value=value).pack(pady=5, padx=20, anchor='w')
    
    # 默认选择第一个策略
    if strategies_list:
        strategy_var.set(strategies_list[0][1])
    
    # 选中的策略
    selected_strategy = None
    
    def on_confirm():
        nonlocal selected_strategy
        selected_strategy = strategy_var.get()
        dialog.destroy()
    
    # 确认按钮
    confirm_btn = ttk.Button(dialog, text="确认", command=on_confirm)
    confirm_btn.pack(pady=10)
    
    # 等待对话框关闭
    dialog.wait_window()
    
    return selected_strategy


def show_backtest_results(parent, backtest_data, metrics, stock_name, strategy_name):
    """显示回测结果"""
    # 创建结果窗口
    result_window = tk.Toplevel(parent)
    result_window.title(f"{stock_name} - 回测结果")
    result_window.geometry("1200x800")
    
    # 创建笔记本（标签页）
    notebook = ttk.Notebook(result_window)
    notebook.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
    
    # 1. 结果指标标签页
    metrics_frame = ttk.Frame(notebook)
    notebook.add(metrics_frame, text="回测指标")
    
    # 创建表格显示指标
    metrics_tree = ttk.Treeview(metrics_frame, show="headings")
    metrics_tree["columns"] = ['指标', '数值']
    
    metrics_tree.column('指标', width=200, anchor='center')
    metrics_tree.column('数值', width=200, anchor='center')
    
    metrics_tree.heading('指标', text='指标', anchor='center')
    metrics_tree.heading('数值', text='数值', anchor='center')
    
    # 添加指标数据
    metrics_data = [
        (f"策略名称 {strategy_name}"),
        ('初始资金', f"{metrics['初始资金']:,.0f} 元"),
        ('最终资产', f"{metrics['最终资产']:,.2f} 元"),
        ('总盈利', f"{metrics['总盈利']:,.2f} 元"),
        ('总收益率', f"{metrics['总收益率']:.2%}"),
        ('总交易日数', f"{metrics['总交易日数']:.0f}"),
        ('市场总收益率', f"{metrics['市场总收益率']:.2%}"),
        ('策略总收益率', f"{metrics['策略总收益率']:.2%}"),
        ('市场年化收益率', f"{metrics['市场年化收益率']:.2%}"),
        ('策略年化收益率', f"{metrics['策略年化收益率']:.2%}"),
        ('市场最大回撤', f"{metrics['市场最大回撤']:.2%}"),
        ('策略最大回撤', f"{metrics['策略最大回撤']:.2%}"),
        ('胜率', f"{metrics['胜率']:.2%}"),
        ('交易次数', f"{metrics['交易次数']:.0f}"),
        ('平均每笔盈利', f"{metrics['平均每笔盈利']:,.2f} 元"),
        ('平均每笔亏损', f"{metrics['平均每笔亏损']:,.2f} 元"),
        ('风险回报比', f"{metrics['风险回报比']:.2f}")
    ]
    
    for item in metrics_data:
        metrics_tree.insert("", tk.END, values=item)
    
    # 滚动条
    metrics_vsb = ttk.Scrollbar(metrics_frame, orient="vertical", command=metrics_tree.yview)
    metrics_tree.configure(yscrollcommand=metrics_vsb.set)
    
    # 布局
    metrics_tree.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=10, pady=10)
    metrics_vsb.pack(side=tk.RIGHT, fill=tk.Y, padx=(0,10), pady=10)
    
    # 2. 收益曲线标签页
    returns_frame = ttk.Frame(notebook)
    notebook.add(returns_frame, text="收益收益曲线")
    
    # 创建图表
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # 绘制累计收益率曲线
    ax.plot(backtest_data['trade_date'], backtest_data['cumulative_market'], label='市场收益', alpha=0.7)
    ax.plot(backtest_data['trade_date'], backtest_data['cumulative_strategy'], label='策略收益', alpha=0.7)
    
    # 添加标题和标签
    ax.set_title(f"{stock_name} - {strategy_name}回测结果", fontsize=14)
    ax.set_xlabel('日期', fontsize=12)
    ax.set_ylabel('累计收益率', fontsize=12)
    ax.legend(fontsize=10)
    ax.grid(True, alpha=0.3)
    
    # 旋转x轴标签
    plt.xticks(rotation=45)
    
    # 调整布局
    plt.tight_layout()
    
    # 嵌入图表
    canvas = FigureCanvasTkAgg(fig, master=returns_frame)
    canvas.draw()
    canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
    
    # 3. 资金曲线标签页
    capital_frame = ttk.Frame(notebook)
    notebook.add(capital_frame, text="资金曲线")
    
    # 创建图表
    fig2, ax2 = plt.subplots(figsize=(10, 6))
    
    # 绘制资金曲线
    ax2.plot(backtest_data['trade_date'], backtest_data['total_asset'], label='总资产', alpha=0.7)
    ax2.axhline(y=metrics['初始资金'], color='r', linestyle='--', label='初始资金', alpha=0.7)
    
    # 添加标题和标签
    ax2.set_title(f"{stock_name} - 资金曲线", fontsize=14)
    ax2.set_xlabel('日期', fontsize=12)
    ax2.set_ylabel('资产价值 (元)', fontsize=12)
    ax2.legend(fontsize=10)
    ax2.grid(True, alpha=0.3)
    
    # 旋转x轴标签
    plt.xticks(rotation=45)
    
    # 调整布局
    plt.tight_layout()
    
    # 嵌入图表
    canvas2 = FigureCanvasTkAgg(fig2, master=capital_frame)
    canvas2.draw()
    canvas2.get_tk_widget().pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
    
    # 4. 交易信号标签页
    signals_frame = ttk.Frame(notebook)
    notebook.add(signals_frame, text="交易信号")
    
    # 创建图表
    fig3, (ax3, ax4) = plt.subplots(2, 1, figsize=(10, 8), gridspec_kw={'height_ratios': [3, 1]})
    
    # 价格走势图
    ax3.plot(backtest_data['trade_date'], backtest_data['close'], label='收盘价', alpha=0.7)
    
    # 绘制移动平均线（如果有）
    if 'short_ma' in backtest_data.columns and 'long_ma' in backtest_data.columns:
        ax3.plot(backtest_data['trade_date'], backtest_data['short_ma'], label='短期均线(20日)', alpha=0.7)
        ax3.plot(backtest_data['trade_date'], backtest_data['long_ma'], label='长期均线(50日)', alpha=0.7)
    
    # 绘制买卖信号
    buy_signals = backtest_data[backtest_data['position'] == 1]
    sell_signals = backtest_data[backtest_data['position'] == -1]
    
    ax3.scatter(buy_signals['trade_date'], buy_signals['close'], marker='^', color='g', label='买入信号', s=100)
    ax3.scatter(sell_signals['trade_date'], sell_signals['close'], marker='v', color='r', label='卖出信号', s=100)
    
    ax3.set_title(f"{stock_name} - 价格走势与交易信号", fontsize=14)
    ax3.set_ylabel('价格', fontsize=12)
    ax3.legend(fontsize=10)
    ax3.grid(True, alpha=0.3)
    
    # 成交量
    ax4.bar(backtest_data['trade_date'], backtest_data['vol'], alpha=0.7)
    ax4.set_title('成交量', fontsize=14)
    ax4.set_xlabel('日期', fontsize=12)
    ax4.set_ylabel('成交量', fontsize=12)
    ax4.grid(True, alpha=0.3)
    
    # 旋转x轴标签
    plt.xticks(rotation=45)
    
    # 调整布局
    plt.tight_layout()
    
    # 嵌入图表
    canvas3 = FigureCanvasTkAgg(fig3, master=signals_frame)
    canvas3.draw()
    canvas3.get_tk_widget().pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
    
    # 5. 交易记录标签页
    trades_frame = ttk.Frame(notebook)
    notebook.add(trades_frame, text="交易记录")
    
    # 筛选有交易的记录
    trade_records = backtest_data[backtest_data['position'] != 0].copy()
    
    if not trade_records.empty:
        # 创建表格显示交易记录
        trades_tree = ttk.Treeview(trades_frame, show="headings")
        trades_tree["columns"] = ['交易日期', '交易类型', '交易价格', '盈亏金额', '收益率', '佣金', '印花税']
        
        # 修改后（添加佣金和印花税列的配置）
        trades_tree.column('交易日期', width=120, anchor='center')
        trades_tree.column('交易类型', width=100, anchor='center')
        trades_tree.column('交易价格', width=100, anchor='center')
        trades_tree.column('盈亏金额', width=120, anchor='center')
        trades_tree.column('收益率', width=100, anchor='center')
        trades_tree.column('佣金', width=100, anchor='center')  # 新增佣金列
        trades_tree.column('印花税', width=100, anchor='center')  # 新增印花税列

        trades_tree.heading('交易日期', text='交易日期', anchor='center')
        trades_tree.heading('交易类型', text='交易类型', anchor='center')
        trades_tree.heading('交易价格', text='交易价格', anchor='center')
        trades_tree.heading('盈亏金额', text='盈亏金额', anchor='center')
        trades_tree.heading('收益率', text='收益率', anchor='center')
        trades_tree.heading('佣金', text='佣金(元)', anchor='center')  # 新增佣金标题
        trades_tree.heading('印花税', text='印花税(元)', anchor='center')  # 新增印花税标题 
        
        # 添加交易数据
        for i, row in trade_records.iterrows():
            trade_type = "买入" if row['position'] == 1 else "卖出"
            trade_price = row['close']
            pnl = row['trade_pnl'] if row['trade_pnl'] != 0 else ""
            pnl_str = f"{pnl:,.2f} 元" if pnl != "" else ""
            return_rate = row['trade_return'] if row['trade_return'] != 0 else ""
            return_rate_str = f"{return_rate:.2%}" if return_rate != "" else ""
            
            # 新增：获取佣金和印花税（保留两位小数）
            commission = f"{row['commission']:.2f}" if row['commission'] > 0 else ""
            stamp_duty = f"{row['stamp_duty']:.2f}" if row['stamp_duty'] > 0 else ""
            
            trades_tree.insert("", tk.END, values=(
                row['trade_date'].strftime("%Y-%m-%d"),
                trade_type,
                f"{trade_price:.2f}",
                pnl_str,
                return_rate_str,
                commission,  # 新增佣金
                stamp_duty   # 新增印花税
            ))
        
        # 滚动条
        trades_vsb = ttk.Scrollbar(trades_frame, orient="vertical", command=trades_tree.yview)
        trades_tree.configure(yscrollcommand=trades_vsb.set)
        
        # 布局
        trades_tree.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=10, pady=10)
        trades_vsb.pack(side=tk.RIGHT, fill=tk.Y, padx=(0,10), pady=10)
    else:
        # 没有交易记录时显示提示
        no_trades_label = ttk.Label(trades_frame, text="在所选日期范围内没有交易记录", font=('Arial', 12))
        no_trades_label.pack(pady=20)


def run_backtest(parent, engine, selected_stock=None):
    """
    运行回测
    """
    try:
        # 1. 选择股票
        if selected_stock is None:
            selected_stock = select_stock_dialog(parent, engine)
            if not selected_stock:
                return
        
        ts_code, stock_name = selected_stock
        
        # 2. 获取日期范围
        min_date, max_date = fetch_date_range(engine)
        if not min_date or not max_date:
            messagebox.showerror("错误", "数据库中没有股票数据")
            return
        
        # 3. 选择回测日期范围
        date_range = select_date_range_dialog(parent, min_date, max_date, stock_name)
        if not date_range:
            return
        
        start_date, end_date = date_range
        
        # 4. 选择策略
        strategy = select_strategy_dialog(parent)
        if not strategy:
            return
        
        # 5. 获取股票数据
        #messagebox.showinfo("提示", "正在获取数据，请稍候...")
        data = fetch_daily_data(engine, ts_code, start_date, end_date)
        
        if data.empty:
            messagebox.showinfo("提示", f"在 {start_date} 至 {end_date} 期间没有找到 {ts_code} 的数据")
            return
        
        # 确保数据按日期排序
        data['trade_date'] = pd.to_datetime(data['trade_date'])
        data = data.sort_values('trade_date')
        
        # 6. 运行回测
        #messagebox.showinfo("提示", "正在运行回测，请稍候...")
        
        # 根据选择的策略运行回测
        strategy_name, strategy_func = strategies.get_strategy_by_name(strategy)
        
        if strategy == 'ma_crossover':
            backtest_data, metrics = strategies.backtest_strategy(
                data, strategy_func, 
                initial_capital=100000,  # 初始资金10万
                short_window=20, 
                long_window=50
            )
        elif strategy == 'price_breakout':
            backtest_data, metrics = strategies.backtest_strategy(
                data, strategy_func,
                initial_capital=100000,  # 初始资金10万
                window=20
            )
        elif strategy == 'rsi_strategy':
            backtest_data, metrics = strategies.backtest_strategy(
                data, strategy_func,
                initial_capital=100000,  # 初始资金10万
                window=14, 
                overbought=70, 
                oversold=30
            )
        elif strategy == 'comprehensive':
            backtest_data, metrics = strategies.backtest_strategy(
                data, strategy_func,
                initial_capital=100000,  # 初始资金10万
                lookback_period=10,
                volume_multiple=2.0,
                short_window=5,
                long_window=20,
                bollinger_window=20,
                num_std=2,
                mean_reversion_window=10,
                z_score_threshold=1.0
            )
        else:
            messagebox.showerror("错误", "未知策略")
            return
        
        # 7. 显示回测结果
        show_backtest_results(parent, backtest_data, metrics, stock_name, strategy_name)
        
    except Exception as e:
        messagebox.showerror("错误", f"回测过程中出错:\n{str(e)}")
        import traceback
        traceback.print_exc()


def display_daily_data(parent, engine, ts_code, stock_name):
    """显示指定股票的日线数据"""
    try:
        data = fetch_daily_data(engine, ts_code)

        if data.empty:
            messagebox.showinfo("提示", f"没有找到 {ts_code} 的日线数据")
            return

        # 创建新窗口
        daily_window = tk.Toplevel(parent)
        daily_window.title(f"{stock_name} ({ts_code}) - 日线数据 (共 {len(data)} 条)")
        daily_window.geometry("1300x800")

        # 主框架
        main_frame = ttk.Frame(daily_window)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # 创建Treeview组件
        tree = ttk.Treeview(main_frame, show="headings")
        tree["columns"] = list(data.columns)

        # 设置列格式
        col_widths = {
            'ts_code': 100,
            'trade_date': 100,
            'open': 80,
            'high': 80,
            'low': 80,
            'close': 80,
            'pre_close': 80,
            'chg_amo': 80,
            'pct_chg': 80,
            'vol': 100,
            'amount': 120
        }

        for col in data.columns:
            width = col_widths.get(col, 100)
            # 修改此处，使数据居中对齐
            tree.column(col, width=width, anchor='center')
            tree.heading(col, text=col, anchor='center')

        # 添加数据
        for i, row in data.iterrows():
            tree.insert("", tk.END, values=list(row))

        # 滚动条
        vsb = ttk.Scrollbar(main_frame, orient="vertical", command=tree.yview)
        hsb = ttk.Scrollbar(main_frame, orient="horizontal", command=tree.xview)
        tree.configure(yscrollcommand=vsb.set, xscrollcommand=hsb.set)

        # 布局
        tree.grid(row=0, column=0, sticky='nsew')
        vsb.grid(row=0, column=1, sticky='ns')
        hsb.grid(row=1, column=0, sticky='ew')

        # 网格配置
        main_frame.grid_rowconfigure(0, weight=1)
        main_frame.grid_columnconfigure(0, weight=1)

        # 添加图表
        if len(data) > 0:
            try:
                fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), gridspec_kw={'height_ratios': [3, 1]})

                # 价格走势图 - 反转数据顺序，使最新数据在右侧
                reversed_data = data.iloc[::-1]
                numeric_cols = ['open', 'high', 'low', 'close']
                for col in numeric_cols:
                    if col in reversed_data.columns:
                        ax1.plot(reversed_data['trade_date'], reversed_data[col], label=col)

                ax1.set_title(f"{stock_name} ({ts_code}) - 价格走势")
                ax1.legend()
                ax1.grid(True)
                plt.setp(ax1.get_xticklabels(), rotation=45)

                # 成交量柱状图 - 反转数据顺序
                if 'vol' in reversed_data.columns:
                    ax2.bar(reversed_data['trade_date'], reversed_data['vol'], width=0.6, color='blue', alpha=0.6)
                    ax2.set_title("成交量")
                    ax2.grid(True)
                    plt.setp(ax2.get_xticklabels(), rotation=45)

                plt.tight_layout()

                # 嵌入图表
                chart_frame = ttk.Frame(daily_window)
                chart_frame.pack(fill=tk.BOTH, expand=True)
                canvas = FigureCanvasTkAgg(fig, master=chart_frame)
                canvas.draw()
                canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

            except Exception as e:
                print(f"创建图表时出错: {e}")

    except Exception as e:
        messagebox.showerror("错误", f"获取日线数据时出错:\n{str(e)}")


def display_stocks_info(engine):
    """显示stocks_info表数据"""
    try:
        info_data = fetch_stocks_info(engine)

        if info_data.empty:
            messagebox.showinfo("提示", "stocks_info表中没有数据")
            return

        # 创建主窗口
        root = tk.Tk()
        root.title("股票基本信息 (双击查看日线数据)")
        root.geometry("1200x800")

        # 主框架
        main_frame = ttk.Frame(root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # 创建Treeview组件
        tree = ttk.Treeview(main_frame, show="headings")
        tree["columns"] = list(info_data.columns)

        # 设置列格式
        col_widths = {
            '序号': 60,  # 新增的序号列
            'ts_code': 80,
            'symbol': 80,
            'name': 80,
            'area': 80,
            'industry': 100,
            'cnspell': 100,
            'market': 60,
            'list_date': 100,
            'st_status': 70,
            'is_hs': 40,
            'act_name': 200,
            'act_ent_type': 80,
            'exchange': 80
        }

        for col in info_data.columns:
            width = col_widths.get(col, 120)
            # 修改此处，使数据居中对齐
            tree.column(col, width=width, anchor='center')
            tree.heading(col, text=col, anchor='center')

        # 添加数据
        for i, row in info_data.iterrows():
            tree.insert("", tk.END, values=list(row))

        # 双击事件
        def on_double_click(event):
            item = tree.selection()[0]
            values = tree.item(item, 'values')
            ts_code = values[1]  # ts_code现在是第二列（因为第一列是序号）
            stock_name = values[3] if len(values) > 3 else ts_code  # name现在是第四列
            display_daily_data(root, engine, ts_code, stock_name)

        tree.bind("<Double-1>", on_double_click)

        # 选中事件
        selected_stock = None
        
        def on_select_item(event):
            nonlocal selected_stock
            if tree.selection():
                item = tree.selection()[0]
                values = tree.item(item, 'values')
                ts_code = values[1]  # ts_code现在是第二列（因为第一列是序号）
                stock_name = values[3] if len(values) > 3 else ts_code  # name现在是第四列
                selected_stock = (ts_code, stock_name)
                
                # 高亮显示选中的行
                for row in tree.get_children():
                    tree.item(row, tags=('selected',) if row == item else ())
                tree.tag_configure('selected', background='lightblue')

        tree.bind("<<TreeviewSelect>>", on_select_item)

        # 滚动条
        vsb = ttk.Scrollbar(main_frame, orient="vertical", command=tree.yview)
        hsb = ttk.Scrollbar(main_frame, orient="horizontal", command=tree.xview)
        tree.configure(yscrollcommand=vsb.set, xscrollcommand=hsb.set)

        # 布局
        tree.grid(row=0, column=0, sticky='nsew')
        vsb.grid(row=0, column=1, sticky='ns')
        hsb.grid(row=1, column=0, sticky='ew')

        # 网格配置
        main_frame.grid_rowconfigure(0, weight=1)
        main_frame.grid_columnconfigure(0, weight=1)

        # 添加功能按钮框架
        button_frame = ttk.Frame(root)
        button_frame.pack(fill=tk.X, padx=10, pady=5)
        
        # 回测按钮
        backtest_btn = ttk.Button(button_frame, text="策略回测", command=lambda: run_backtest(root, engine, selected_stock))
        backtest_btn.pack(side=tk.LEFT, padx=5)
        
        # 添加搜索功能
        search_frame = ttk.Frame(root)
        search_frame.pack(fill=tk.X, padx=10, pady=5)

        ttk.Label(search_frame, text="搜索:").pack(side=tk.LEFT, padx=5)
        search_var = tk.StringVar()
        search_entry = ttk.Entry(search_frame, textvariable=search_var)
        search_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)

        def update_tree():
            search_term = search_var.get().lower()
            for item in tree.get_children():
                values = tree.item(item, 'values')
                if any(search_term in str(v).lower() for v in values):
                    tree.item(item, tags=('match',))
                    tree.selection_set(item)
                else:
                    tree.item(item, tags=('no_match',))

            tree.tag_configure('match', background='yellow')
            tree.tag_configure('no_match', background='white')

        search_var.trace_add("write", lambda *args: update_tree())

        try:
            root.mainloop()
        except KeyboardInterrupt:
            print("程序被用户中断")
        except Exception as e:
            messagebox.showerror("错误", f"程序运行出错:\n{str(e)}")

    except Exception as e:
        messagebox.showerror("错误", f"获取股票信息时出错:\n{str(e)}")


def main():
    """主函数"""
    try:
        # 创建数据库连接
        engine = create_engine(f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}?charset={DB_CONFIG['charset']}")
        
        # 测试数据库连接
        with engine.connect():
            print("数据库连接成功")
            
        # 显示股票信息
        display_stocks_info(engine)
        
    except Exception as e:
        messagebox.showerror("错误", f"数据库连接失败:\n{str(e)}")


if __name__ == "__main__":
    main()
